package main

import (
	"fmt"
	"sync"
	"sync/atomic"
	"time"
)

// RateLimiter 限流器结构体
type RateLimiter struct {
	limit   int32      // 最大允许并发请求数
	current int32      // 当前请求数
	mu      sync.Mutex // Cond 所需的锁
	cond    *sync.Cond // 条件变量用于阻塞/唤醒
	closed  int32      // 是否关闭
}

// NewRateLimiter 创建一个新的限流器
func NewRateLimiter(limit int32) *RateLimiter {
	if limit <= 0 {
		panic("limit must be greater than 0")
	}
	limiter := &RateLimiter{
		limit:   limit,
		current: 0,
	}
	limiter.cond = sync.NewCond(&limiter.mu)
	return limiter
}

// Acquire 获取一个请求许可
func (r *RateLimiter) Acquire() {
	if atomic.LoadInt32(&r.closed) == 1 {
		panic("rate limiter is closed")
	}

	r.mu.Lock()
	defer r.mu.Unlock()

	for atomic.LoadInt32(&r.current) >= r.limit {
		r.cond.Wait() // 阻塞直到有空位
	}
	atomic.AddInt32(&r.current, 1)
}

// Release 释放一个请求许可
func (r *RateLimiter) Release() {
	if atomic.LoadInt32(&r.closed) == 1 {
		panic("rate limiter is closed")
	}

	newCount := atomic.AddInt32(&r.current, -1)
	r.mu.Lock()
	defer r.mu.Unlock()
	if newCount < r.limit {
		r.cond.Signal() // 唤醒一个等待者
	}
}

// Close 关闭限流器，后续调用 Acquire 会 panic
func (r *RateLimiter) Close() {
	atomic.StoreInt32(&r.closed, 1)
	r.mu.Lock()
	r.cond.Broadcast() // 唤醒所有等待中的 goroutine，让它们退出
	r.mu.Unlock()
}

// ---- 测试代码 ----

func worker(id int, limiter *RateLimiter, wg *sync.WaitGroup) {
	defer wg.Done()

	fmt.Printf("Worker %d is waiting to acquire...\n", id)
	limiter.Acquire()
	fmt.Printf("Worker %d acquired the permit.\n", id)

	time.Sleep(500 * time.Millisecond) // 模拟处理时间

	limiter.Release()
	fmt.Printf("Worker %d released the permit.\n", id)
}

func main() {
	limiter := NewRateLimiter(3) // 同时最多允许 3 个并发请求
	var wg sync.WaitGroup

	for i := 1; i <= 10; i++ {
		wg.Add(1)
		go worker(i, limiter, &wg)
		time.Sleep(100 * time.Millisecond) // 控制启动频率
	}

	wg.Wait()
	fmt.Println("All workers done.")

	limiter.Close()
}
